# Copyright (c) 2025 Deep Robotics  # 2025年Deep Robotics公司版权所有
# SPDX-License-Identifier: BSD 3-Clause  # 遵循BSD 3-Clause许可证

# Copyright (c) 2024-2025 Ziqi Fan  # 2024-2025年Ziqi Fan版权所有
# SPDX-License-Identifier: Apache-2.0  # 遵循Apache-2.0许可证

# 导入isaaclab的configclass装饰器
from isaaclab.utils import configclass

# 导入粗糙地形运动速度环境配置基类
from rl_training.tasks.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg
# # 可选：导入射线传感器的网格模式配置
##
# Pre-defined configs
##
# 导入机器人配置
from rl_training.assets.deeprobotics import DEEPROBOTICS_LITE3_CFG  # isort: skip


@configclass
class DeeproboticsLite3RoughEnvCfg(LocomotionVelocityRoughEnvCfg):
    base_link_name = "TORSO"  # 机器人主躯干链接名
    foot_link_name = ".*_FOOT"  # 足部链接名的正则表达式
    # fmt: off
    joint_names = [  # 关节名称列表
        "FL_HipX_joint", "FL_HipY_joint", "FL_Knee_joint",
        "FR_HipX_joint", "FR_HipY_joint", "FR_Knee_joint",
        "HL_HipX_joint", "HL_HipY_joint", "HL_Knee_joint",
        "HR_HipX_joint", "HR_HipY_joint", "HR_Knee_joint",
    ]

    link_names = [  # 机器人所有链接名称
       'TORSO', 
       'FL_HIP', 'FR_HIP', 'HL_HIP', 'HR_HIP', 
       'FL_THIGH', 'FR_THIGH', 'HL_THIGH', 'HR_THIGH', 
       'FL_SHANK', 'FR_SHANK', 'HL_SHANK', 'HR_SHANK', 
       'FL_FOOT', 'FR_FOOT', 'HL_FOOT', 'HR_FOOT',
    ]
    # fmt: on

    def __post_init__(self):
        # 调用父类的post_init方法
        super().__post_init__()

        # ------------------------------Sence------------------------------
        self.scene.robot = DEEPROBOTICS_LITE3_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")  # 设置机器人模型路径
        self.scene.height_scanner.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name  # 高度扫描仪主路径
        self.scene.height_scanner_base.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name  # 高度扫描仪基准路径
        self.scene.height_scanner.pattern_cfg.resolution = 0.07 # 设置高度扫描仪分辨率

        # ------------------------------Observations------------------------------
        self.observations.policy.base_lin_vel.scale = 2.0
        self.observations.policy.base_ang_vel.scale = 0.25  # 基础角速度缩放
        self.observations.policy.joint_pos.scale = 1.0      # 关节位置缩放
        self.observations.policy.joint_vel.scale = 0.05     # 关节速度缩放
        self.observations.policy.base_lin_vel = None        # 不观测基础线速度
        self.observations.policy.height_scan = None         # 不观测高度扫描
        self.observations.policy.joint_pos.params["asset_cfg"].joint_names = self.joint_names  # 关节位置观测关节名
        self.observations.policy.joint_vel.params["asset_cfg"].joint_names = self.joint_names  # 关节速度观测关节名

        # ------------------------------Actions------------------------------
        # 降低动作缩放比例
        self.actions.joint_pos.scale = {".*_HipX_joint": 0.125, "^(?!.*_HipX_joint).*": 0.25}  # HipX关节动作缩放更小
        self.actions.joint_pos.clip = {".*": (-100.0, 100.0)}  # 所有关节动作裁剪范围
        self.actions.joint_pos.joint_names = self.joint_names  # 动作关节名

        # ------------------------------Events------------------------------
        # 随机重置基座参数
        self.events.randomize_reset_base.params = {
            "pose_range": {
                "x": (-1.0, 1.0),
                "y": (-1.0, 1.0),
                "z": (0.0, 0.0),
                "roll": (-0.3, 0.3),
                "pitch": (-0.3, 0.3),
                "yaw": (-3.14, 3.14),
            },
            "velocity_range": {
                "x": (-0.2, 0.2),
                "y": (-0.2, 0.2),
                "z": (-0.2, 0.2),
                "roll": (-0.05, 0.05),
                "pitch": (-0.05, 0.05),
                "yaw": (-0.0, 0.0),
            },
        }
        self.events.randomize_rigid_body_mass.params["asset_cfg"].body_names = self.link_names # 随机刚体质量的链接名
        self.events.randomize_rigid_body_mass_base = None  # 不随机基座质量
        self.events.randomize_com_positions.params["asset_cfg"].body_names = self.base_link_name # 随机质心位置的链接名
        # self.events.randomize_com_positions = None  # 可选：不随机质心位置
        self.events.randomize_apply_external_force_torque = None  # 不施加外力
        self.events.randomize_push_robot = None  # 不推机器人
        self.events.randomize_actuator_gains.params["asset_cfg"].joint_names = self.joint_names  # 随机执行器增益的关节名

        # 缩小地形尺寸以适应小型机器人
        self.scene.terrain.terrain_generator.sub_terrains["boxes"].grid_height_range = (0.025, 0.1)  # box地形高度范围
        self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_range = (0.01, 0.06)  # 随机粗糙地形噪声范围
        self.scene.terrain.terrain_generator.sub_terrains["random_rough"].noise_step = 0.01  # 噪声步长

        # ------------------------------Rewards------------------------------

        # Root penalties
        # Z轴方向线速度和XY方向角速度惩罚，不希望机器人在Z轴上下运动，或者在XY平面有过多旋转，期望平稳移动
        self.rewards.lin_vel_z_l2.weight = -2.0    # z方向线速度惩罚
        self.rewards.ang_vel_xy_l2.weight = -0.05  # xy方向角速度惩罚
        # 基座在Z轴方向保持平稳，不希望它翻滚或侧翻
        self.rewards.flat_orientation_l2.weight = -5.0  # 姿态平坦度惩罚
        self.rewards.base_height_l2.weight = -10.0      # 基座高度惩罚
        self.rewards.base_height_l2.params["target_height"] = 0.35  # 目标高度
        self.rewards.base_height_l2.params["asset_cfg"].body_names = [self.base_link_name]  # 作用链接

        # Joint penalties
        # 不期望超出扭矩限制，超出了会烧坏电机
        self.rewards.joint_torques_l2.weight = -2.5e-5  # 关节力矩惩罚
        # 关节加速度，太大会导致电机温度升高进而烧坏电机
        self.rewards.joint_acc_l2.weight = -1e-8  # 关节加速度惩罚
        # 节省电量
        self.rewards.joint_power.weight = -2e-5  # 关节功率惩罚
        # 这个动作和下一个动作之间的差异，我们不希望动作变化太大，否则机器人会不稳定，我们将惩罚这种动作
        self.rewards.action_rate_l2.weight = -0.02  # 动作变化率惩罚
        # self.rewards.smoothness_2.weight = -0.0075  # 可选：动作平滑度惩罚

        # Contact sensor不期望它在地上翻滚
        self.rewards.undesired_contacts.weight = -0.5  # 非期望接触惩罚
        self.rewards.undesired_contacts.params["sensor_cfg"].body_names = [f"^(?!.*{self.foot_link_name}).*"]
        # 希望机器人足端轻触地面，接触力越小越好。
        self.rewards.contact_forces.weight = -2e-2  # 接触力惩罚
        self.rewards.contact_forces.params["sensor_cfg"].body_names = [self.foot_link_name]
        # Velocity-tracking rewards
        self.rewards.track_lin_vel_xy_exp.weight = 1.2  # xy方向速度跟踪奖励
        self.rewards.track_ang_vel_z_exp.weight = 0.6   # z方向角速度跟踪奖励
        # Others
        self.rewards.feet_air_time.weight = 1.0  # 足部离地时间奖励
        self.rewards.feet_air_time.params["threshold"] = 0.5  # 离地阈值
        self.rewards.feet_air_time.params["sensor_cfg"].body_names = [self.foot_link_name]  # 足部链接
        self.rewards.feet_air_time_variance.weight = -8.0  # 足部离地时间方差惩罚
        self.rewards.feet_air_time_variance.params["sensor_cfg"].body_names = [self.foot_link_name]
        # 机器人踩在地上时，可能会滑动，使用feet_slide惩罚滑动
        self.rewards.feet_slide.weight = -0.05  # 足部滑动惩罚
        self.rewards.feet_slide.params["sensor_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_slide.params["asset_cfg"].body_names = [self.foot_link_name]
        # 踏步时，实际想要的足端高度
        self.rewards.feet_height.weight = -0.2  # 足部高度惩罚
        self.rewards.feet_height.params["target_height"] = 0.05
        self.rewards.feet_height.params["asset_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_height_body.weight = -2.5  # 足部相对身体高度惩罚
        self.rewards.feet_height_body.params["target_height"] = -0.35
        self.rewards.feet_height_body.params["asset_cfg"].body_names = [self.foot_link_name]
        # 没有任何指令时，希望机器人完全静止，使用stand_still函数，作为关节偏差与默认关节位姿比较，如果偏离就惩罚
        self.rewards.stand_still.weight = -0.3  # 静止惩罚
        self.rewards.stand_still.params["asset_cfg"].joint_names = self.joint_names
        self.rewards.stand_still.params["command_threshold"] = 0.1  # 命令阈值
        # 期望腿垂直踏在地面
        self.rewards.joint_deviation_l1.weight = -0.5  # 关节偏差惩罚
        self.rewards.joint_deviation_l1.params["asset_cfg"].joint_names = [".*HipX.*"]


        # 如果奖励权重为0，则禁用该奖励
        if self.__class__.__name__ == "DeeproboticsLite3RoughEnvCfg":
            self.disable_zero_weight_rewards()

        # ------------------------------Terminations------------------------------
        self.terminations.illegal_contact = None      # 非法接触终止条件
        # self.terminations.bad_orientation_2 = None  # 可选：姿态异常终止条件

        # ------------------------------Curriculums------------------------------
        # self.curriculum.command_levels.params["range_multiplier"] = (0.2, 1.0)  # 可选：命令难度课程
        self.curriculum.command_levels = None  # 不使用命令课程

        # ------------------------------Commands------------------------------
        self.commands.base_velocity.ranges.lin_vel_x = (-1.5, 1.5)  # x方向速度命令范围
        self.commands.base_velocity.ranges.lin_vel_y = (-0.8, 0.8)  # y方向速度命令范围
        self.commands.base_velocity.ranges.ang_vel_z = (-1.5, 1.5)  # z方向角速度命令范围